# ocr_processor.py
import os
import re
import logging
import time
from typing import List, Dict
from PIL import Image, ImageEnhance, ImageOps
from aip import AipOcr
import pandas as pd

# 日志配置
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s: %(message)s',
    handlers=[
        logging.FileHandler('ocr_processor.log'),
        logging.StreamHandler()
    ]
)

class OCRProcessor:
    def __init__(self, app_id, api_key, secret_key):
        self.client = AipOcr(app_id, api_key, secret_key)
        self.surnames = self._load_surnames()
        self.amount_pattern = r'[¥￥$]?\s*\d{1,3}(?:,\d{3})*(?:\.\d+)?\s*元?|\d+(?:\.\d+)?\s*[元圆整]|(?:[壹贰叁肆伍陆柒捌玖拾佰仟萬]+)圆?整?'
        
        # 配置校验
        if not all([app_id, api_key, secret_key]):
            logging.error("百度OCR配置不完整，请检查APP_ID/API_KEY/SECRET_KEY")
            raise ValueError("Invalid Baidu OCR Configuration")

    def _load_surnames(self) -> List[str]:
        """加载扩展姓氏库"""
        return [
            '王','李','张','刘','陈','杨','黄','赵','吴','周','徐','孙','马','朱',
            '胡','郭','何','高','林','罗','郑','梁','谢','宋','唐','许','韩','冯',
            '邓','曹','彭','曾','肖','田','董','袁','潘','于','蒋','蔡','余','杜',
            '叶','程','苏','魏','吕','丁','任','沈','姚','卢','姜','崔','钟','谭',
            '陆','汪','范','金','石','廖','贾','夏','韦','付','方','白','邹','孟',
            '熊','秦','邱','江','尹','薛','闫','段','雷','侯','龙','史','陶','黎',
            '贺','顾','毛','郝','龚','邵','欧阳','司马','上官','诸葛','夏侯'  # 补充复姓
        ]

    def _preprocess_image(self, image_path: str) -> Image.Image:
        """增强图像预处理"""
        try:
            img = Image.open(image_path)
            
            # 自动旋转校正
            img = ImageOps.exif_transpose(img)
            
            # 图像增强
            img = img.convert('L')  # 转灰度
            enhancer = ImageEnhance.Contrast(img)
            img = enhancer.enhance(2.0)  # 提高对比度
            
            # 二值化处理
            threshold = 150
            img = img.point(lambda p: 255 if p > threshold else 0)
            
            return img
        except Exception as e:
            logging.error(f"图像预处理失败: {str(e)}")
            raise

    def _ocr_retry(self, image_data: bytes, retries=3) -> dict:
        """带重试机制的OCR请求"""
        for i in range(retries):
            try:
                result = self.client.basicAccurate(image_data, {
                    'language_type': 'CHN_ENG',
                    'detect_direction': 'true'
                })
                if 'words_result' in result:
                    return result
                logging.warning(f"OCR响应异常，重试 {i+1}/{retries}")
            except Exception as e:
                logging.error(f"OCR请求失败: {str(e)}")
            time.sleep(1)
        return {'words_result': []}

    def _classify_text(self, text: str) -> Dict:
        """改进的分类逻辑"""
        # 金额识别（支持中文大写）
        amounts = []
        for match in re.finditer(self.amount_pattern, text):
            amount = match.group().replace(' ', '')
            amounts.append(amount)
        
        # 姓名识别（放宽长度限制）
        names = []
        for match in re.finditer(r'[\u4e00-\u9fa5]{2,4}', text):
            name = match.group()
            if name[0] in self.surnames or (len(name)==2 and name[-1] in ['某']):
                names.append(name)
        
        return {'names': names, 'amounts': amounts}

    def _process_image(self, image_path: str) -> List[Dict]:
        """处理单张图片"""
        try:
            # 图像预处理
            img = self._preprocess_image(image_path)
            img.save("temp_ocr.jpg")
            
            with open("temp_ocr.jpg", "rb") as f:
                result = self._ocr_retry(f.read())
            
            # 合并识别文本
            text = ' '.join([w['words'] for w in result.get('words_result', [])])
            classified = self._classify_text(text)
            
            # 智能对齐数据
            paired_data = []
            min_len = min(len(classified['names']), len(classified['amounts']))
            for i in range(min_len):
                paired_data.append({
                    '姓名': classified['names'][i],
                    '金额': classified['amounts'][i]
                })
            
            # 处理剩余未配对数据
            if len(classified['names']) > min_len:
                for name in classified['names'][min_len:]:
                    paired_data.append({'姓名': name, '金额': None})
            elif len(classified['amounts']) > min_len:
                for amount in classified['amounts'][min_len:]:
                    paired_data.append({'姓名': None, '金额': amount})
            
            return paired_data
        except Exception as e:
            logging.error(f"处理图片失败: {image_path} - {str(e)}")
            return []
        finally:
            if os.path.exists("temp_ocr.jpg"):
                os.remove("temp_ocr.jpg")

    def process_folder(self, input_folder: str, output_excel: str) -> bool:
        """处理整个文件夹"""
        # 输入验证
        if not os.path.exists(input_folder):
            logging.error(f"输入文件夹不存在: {input_folder}")
            return False
            
        valid_files = [
            f for f in os.listdir(input_folder)
            if f.lower().endswith(('.png','.jpg','.jpeg','.bmp'))
        ]
        if not valid_files:
            logging.error("输入文件夹中没有支持的图片文件（需.jpg/.png等）")
            return False

        all_data = []
        for filename in valid_files:
            image_path = os.path.join(input_folder, filename)
            rows = self._process_image(image_path)
            
            for row in rows:
                if row['姓名'] or row['金额']:
                    all_data.append({
                        '文件名': filename,
                        '姓名': row['姓名'] or '未识别',
                        '金额': row['金额'] or '未识别'
                    })
        
        # 生成结果
        if all_data:
            df = pd.DataFrame(all_data)
            # 智能过滤空值
            df = df[(df['姓名'] != '未识别') | (df['金额'] != '未识别')]
            
            if not df.empty:
                df.to_excel(output_excel, index=False)
                logging.info(f"成功生成文件: {output_excel}，共{len(df)}条记录")
                return True
        
        logging.warning("未找到有效数据，请检查：\n1. 图片是否包含中文姓名和金额\n2. 图片清晰度是否足够\n3. OCR配置是否正确")
        return False

if __name__ == "__main__":
    # 配置信息（需替换为实际值）
    CONFIG = {
        'APP_ID': '',
        'API_KEY': '',
        'SECRET_KEY': '',
        'INPUT_FOLDER': './images',
        'OUTPUT_EXCEL': 'result.xlsx'
    }
    
    try:
        processor = OCRProcessor(
            CONFIG['APP_ID'],
            CONFIG['API_KEY'],
            CONFIG['SECRET_KEY']
        )
        success = processor.process_folder(
            CONFIG['INPUT_FOLDER'],
            CONFIG['OUTPUT_EXCEL']
        )
        if not success:
            logging.error("处理失败，请查看日志文件 ocr_processor.log")
    except Exception as e:
        logging.error(f"系统错误: {str(e)}")